#include "usbd_core.h"
#include "usbd_msc.h"

#define MSC_IN_EP 0x81
#define MSC_OUT_EP 0x02

#define USBD_VID 0x36b7
#define USBD_PID 0xFFFF
#define USBD_MAX_POWER 100
#define USBD_LANGID_STRING 1033

#define USB_CONFIG_SIZE (9 + MSC_DESCRIPTOR_LEN)

//#define USE_INTERNAL_FLASH
//#define USE_SD_NAND
#define USE_SPI_NAND_FLASH
//#define USE_SPI_NOR_FLASH

const uint8_t msc_flash_descriptor[] = {
    USB_DEVICE_DESCRIPTOR_INIT(USB_2_0, 0x00, 0x00, 0x00, USBD_VID, USBD_PID, 0x0200, 0x01),
    USB_CONFIG_DESCRIPTOR_INIT(USB_CONFIG_SIZE, 0x01, 0x01, USB_CONFIG_BUS_POWERED, USBD_MAX_POWER),
    MSC_DESCRIPTOR_INIT(0x00, MSC_OUT_EP, MSC_IN_EP, 0x02),
    ///////////////////////////////////////
    /// string0 descriptor
    ///////////////////////////////////////
    USB_LANGID_INIT(USBD_LANGID_STRING),
    ///////////////////////////////////////
    /// string1 descriptor
    ///////////////////////////////////////
    0x0A,                       /* bLength */
    USB_DESCRIPTOR_TYPE_STRING, /* bDescriptorType */
    'P', 0x00,                  /* wcChar0 */
    'U', 0x00,                  /* wcChar1 */
    'Y', 0x00,                  /* wcChar2 */
    'A', 0x00,                  /* wcChar3 */
    ///////////////////////////////////////
    /// string2 descriptor
    ///////////////////////////////////////
    0x1C,                       /* bLength */
    USB_DESCRIPTOR_TYPE_STRING, /* bDescriptorType */
    'P', 0x00,                  /* wcChar0 */
    'U', 0x00,                  /* wcChar1 */
    'Y', 0x00,                  /* wcChar2 */
    'A', 0x00,                  /* wcChar3 */
    ' ', 0x00,                  /* wcChar4 */
    'M', 0x00,                  /* wcChar5 */
    'S', 0x00,                  /* wcChar6 */
    'C', 0x00,                  /* wcChar7 */
    ' ', 0x00,                  /* wcChar8 */
    'D', 0x00,                  /* wcChar9 */
    'E', 0x00,                  /* wcChar10 */
    'M', 0x00,                  /* wcChar11 */
    'O', 0x00,                  /* wcChar12 */
    ///////////////////////////////////////
    /// string3 descriptor
    ///////////////////////////////////////
    0x16,                       /* bLength */
    USB_DESCRIPTOR_TYPE_STRING, /* bDescriptorType */
    '2', 0x00,                  /* wcChar0 */
    '0', 0x00,                  /* wcChar1 */
    '2', 0x00,                  /* wcChar2 */
    '2', 0x00,                  /* wcChar3 */
    '1', 0x00,                  /* wcChar4 */
    '2', 0x00,                  /* wcChar5 */
    '3', 0x00,                  /* wcChar6 */
    '4', 0x00,                  /* wcChar7 */
    '5', 0x00,                  /* wcChar8 */
    '6', 0x00,                  /* wcChar9 */
#ifdef CONFIG_USB_HS
    ///////////////////////////////////////
    /// device qualifier descriptor
    ///////////////////////////////////////
    0x0a, USB_DESCRIPTOR_TYPE_DEVICE_QUALIFIER, 0x00, 0x02, 0x00, 0x00, 0x00, 0x40, 0x01, 0x00,
#endif
    0x00};

void usbd_configure_done_callback(void)
{ /* do nothing */
}

#if defined(USE_INTERNAL_FLASH)

#define DISK_BASE_ADDR 0x08010000
#define DISK_PAGE_SIZE FLASH_PAGE_SIZE
#define BLOCK_SIZE 512
#define BLOCK_COUNT (0x8060000 - DISK_BASE_ADDR) / BLOCK_SIZE

void usbd_msc_get_cap(uint8_t lun, uint32_t* block_num, uint16_t* block_size)
{
    *block_num = BLOCK_COUNT;
    *block_size = BLOCK_SIZE;
}

int usbd_msc_sector_read(uint32_t sector, uint8_t* buffer, uint32_t length)
{
    uint32_t *pSrc, *pDst;
    uint32_t  index = 0;
    pSrc = (uint32_t*)(DISK_BASE_ADDR + sector * BLOCK_SIZE);
    pDst = (uint32_t*)buffer;

    for (index = 0; index < (length / 4); index++)
    {
        pDst[index] = pSrc[index];
    }
    return 0;
}

int usbd_msc_sector_write(uint32_t sector, uint8_t* buffer, uint32_t length)
{
    uint32_t  index;
    uint32_t* ptr;
    uint32_t  start_addr = DISK_BASE_ADDR + sector * BLOCK_SIZE;
    uint32_t  end_addr = start_addr + length;
    uint32_t  current_addr = start_addr;

    LL_FLASH_Unlock(FLASH); /* unlock FLASH */

    /* erase */
    while (current_addr < end_addr)
    {
        while (LL_FLASH_IsActiveFlag_BUSY(FLASH) == 1) /* Wait Busy=0 */
            ;
        LL_FLASH_EnableIT_EOP(FLASH);                  /* Enable EOP */
        LL_FLASH_EnablePageErase(FLASH);               /* Enable Page Erase */
        LL_FLASH_SetEraseAddress(FLASH, current_addr); /* Set Erase Address */
        while (LL_FLASH_IsActiveFlag_BUSY(FLASH) == 1) /* Wait Busy=0 */
            ;
        while (LL_FLASH_IsActiveFlag_EOP(FLASH) == 0) /* Wait EOP=1 */
            ;
        LL_FLASH_ClearFlag_EOP(FLASH);    /* Clear EOP */
        LL_FLASH_DisableIT_EOP(FLASH);    /* Disable EOP */
        LL_FLASH_DisablePageErase(FLASH); /* Disable Page Erase */
        current_addr += DISK_PAGE_SIZE;
    }

    /* write */
    current_addr = start_addr;
    for (index = 0; index < (length + DISK_PAGE_SIZE - 1) / DISK_PAGE_SIZE; index++)
    {
        ptr = (uint32_t*)(buffer + DISK_PAGE_SIZE * index);
        while (LL_FLASH_IsActiveFlag_BUSY(FLASH) == 1) /* Wait Busy=0 */
            ;
        LL_FLASH_EnableIT_EOP(FLASH);                   /* Enable EOP */
        LL_FLASH_EnablePageProgram(FLASH);              /* Enable Program */
        LL_FLASH_PageProgram(FLASH, current_addr, ptr); /* Page Program */
        while (LL_FLASH_IsActiveFlag_BUSY(FLASH) == 1)  /* Wait Busy=0 */
            ;
        while (LL_FLASH_IsActiveFlag_EOP(FLASH) == 0) /* Wait EOP=1 */
            ;
        LL_FLASH_ClearFlag_EOP(FLASH);      /* Clear EOP */
        LL_FLASH_DisableIT_EOP(FLASH);      /* Disable EOP */
        LL_FLASH_DisablePageProgram(FLASH); /* Disable Program */
        current_addr += DISK_PAGE_SIZE;
    }

    LL_FLASH_Lock(FLASH); /* Lock FLASH */
    return 0;
}
#endif



#if defined(USE_SPI_NAND_FLASH)

#include "flashDriver.h"
uint32_t flashID = 0;
uint8_t  protReg;
uint8_t  confReg;
uint8_t  statReg;

#define BLOCK_SIZE 512
#define BLOCK_COUNT 262144 // 128MB = 262144 sectors (512B each)
#define NAND_PAGE_SIZE 2048
#define SECTORS_PER_PAGE (NAND_PAGE_SIZE / BLOCK_SIZE) // 4 sectors per page
#define NAND_BLOCK_SIZE (64 * NAND_PAGE_SIZE) // 128KB per block

static uint8_t page_buffer[NAND_PAGE_SIZE];

void usbd_msc_get_cap(uint8_t lun, uint32_t *block_num, uint16_t *block_size)
{
    *block_size = BLOCK_SIZE;
    *block_num = BLOCK_COUNT;
}

int usbd_msc_sector_read(uint32_t sector, uint8_t *buffer, uint32_t length)
{
    if (sector >= BLOCK_COUNT || length != BLOCK_SIZE) {
        rt_kprintf("Read error: sector %lu, length %lu\n", sector, length);
        return -1;
    }
    uint32_t page = sector / SECTORS_PER_PAGE;
    uint16_t offset = (sector % SECTORS_PER_PAGE) * BLOCK_SIZE;

    // Read page
    if (readFlashPage(page_buffer, page * NAND_PAGE_SIZE, NAND_PAGE_SIZE) != HAL_OK) {
        rt_kprintf("Read page %lu failed\n", page);
        return -1;
    }
    memcpy(buffer, &page_buffer[offset], BLOCK_SIZE);
    rt_kprintf("Read sector %lu OK, first 4 bytes: %02X %02X %02X %02X\n",
               sector, buffer[0], buffer[1], buffer[2], buffer[3]);
    return 0;
}

int usbd_msc_sector_write(uint32_t sector, uint8_t *buffer, uint32_t length)
{
    if (sector >= BLOCK_COUNT || length % BLOCK_SIZE != 0) {
        rt_kprintf("Write error: sector %lu, length %lu\n", sector, length);
        return -1;
    }

    uint32_t num_sectors = length / BLOCK_SIZE;
    for (uint32_t i = 0; i < num_sectors; i++) {
        uint32_t curr_sector = sector + i;
        uint8_t *curr_buffer = buffer + (i * BLOCK_SIZE);
        uint32_t page = curr_sector / SECTORS_PER_PAGE;
        uint16_t offset = (curr_sector % SECTORS_PER_PAGE) * BLOCK_SIZE;
        uint32_t block = page / 64;

        rt_kprintf("Writing sector %lu, first 4 bytes: %02X %02X %02X %02X\n",
                   curr_sector, curr_buffer[0], curr_buffer[1], curr_buffer[2], curr_buffer[3]);

        // Erase block if first sector in block
        if (curr_sector % (NAND_BLOCK_SIZE / BLOCK_SIZE) == 0) {
            rt_kprintf("Erasing block %lu\n", block);
            // Retry erase up to 3 times
            for (uint8_t retry = 0; retry < 3; retry++) {
                blockErase(block * NAND_BLOCK_SIZE);
                waitForNotBusy();
                // Add delay for erase completion
                for (volatile uint32_t d = 0; d < 3000; d++); // ~3ms delay
                uint8_t status = checkStatusReg(STATUS_REG_ADDRESS);
                if (status & 0x20) {
                    rt_kprintf("Erase failed at block %lu, retry %u, status: 0x%02X\n",
                               block, retry + 1, status);
                    continue;
                }
                // Verify erase
                if (readFlashPage(page_buffer, page * NAND_PAGE_SIZE, NAND_PAGE_SIZE) != HAL_OK) {
                    rt_kprintf("Read page %lu failed after erase\n", page);
                    return -1;
                }
                bool erased = true;
                for (uint32_t j = 0; j < NAND_PAGE_SIZE; j++) {
                    if (page_buffer[j] != 0xFF) {
                        rt_kprintf("Erase verification failed at page %lu, byte %lu: %02X\n",
                                   page, j, page_buffer[j]);
                        erased = false;
                        break;
                    }
                }
                if (erased) {
                    rt_kprintf("Block %lu erased successfully\n", block);
                    break;
                }
                if (retry == 2) {
                    rt_kprintf("Erase failed after 3 retries at block %lu\n", block);
                    return -1;
                }
            }
        }

        // Initialize page buffer to 0xFF
        memset(page_buffer, 0xFF, NAND_PAGE_SIZE);
        // Read current page to preserve other sectors
        if (readFlashPage(page_buffer, page * NAND_PAGE_SIZE, NAND_PAGE_SIZE) != HAL_OK) {
            rt_kprintf("Read page %lu failed\n", page);
            return -1;
        }

        // Update page buffer
        memcpy(&page_buffer[offset], curr_buffer, BLOCK_SIZE);

        // Enable write and verify WEL
        enableWriteFlash();
        waitForNotBusy();
        uint8_t status = checkStatusReg(STATUS_REG_ADDRESS);
        if (!(status & 0x02)) {
            rt_kprintf("WEL not set for page %lu, status: 0x%02X\n", page, status);
            return -1;
        }

        // Write page
        if (writeFlashPage(page_buffer, page * NAND_PAGE_SIZE, NAND_PAGE_SIZE) != HAL_OK) {
            rt_kprintf("Write page %lu failed\n", page);
            return -1;
        }
        waitForNotBusy();
        status = checkStatusReg(STATUS_REG_ADDRESS);
        if (status & 0x10) {
            rt_kprintf("Write failed at page %lu, status: 0x%02X\n", page, status);
            return -1;
        }

        // Verify write
        if (readFlashPage(page_buffer, page * NAND_PAGE_SIZE, NAND_PAGE_SIZE) != HAL_OK) {
            rt_kprintf("Verify read page %lu failed\n", page);
            return -1;
        }
        uint8_t verify_buffer[BLOCK_SIZE];
        memcpy(verify_buffer, &page_buffer[offset], BLOCK_SIZE);
        if (memcmp(verify_buffer, curr_buffer, BLOCK_SIZE) != 0) {
            rt_kprintf("Verification failed at sector %lu, first 4 bytes: wrote %02X %02X %02X %02X, read %02X %02X %02X %02X\n",
                       curr_sector,
                       curr_buffer[0], curr_buffer[1], curr_buffer[2], curr_buffer[3],
                       verify_buffer[0], verify_buffer[1], verify_buffer[2], verify_buffer[3]);
            return -1;
        }

        rt_kprintf("Write sector %lu OK\n", curr_sector);
    }

    return 0;
}

#endif

struct usbd_interface intf0;

void msc_ram_init(void)
{

#if defined(USE_SPI_NAND_FLASH)
    initFlash();
#endif
    
    usbd_desc_register(msc_flash_descriptor);
    usbd_add_interface(usbd_msc_init_intf(&intf0, MSC_OUT_EP, MSC_IN_EP));

    usbd_initialize();
}
